from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional

app = FastAPI()

class QueryRequest(BaseModel):
    query: str
    top_k: int = 3
    modes: List[str] = ["q-c"]

class SearchResult(BaseModel):
    qc_results: List[dict]
    es_results: List[dict]
    gpt_answer: str

@app.post("/api/query", response_model=SearchResult)
async def process_query(request: QueryRequest):
    try:
        # 将单个查询转换为列表格式
        query_list = [request.query]
        
        # 调用你的查询处理器
        q_q, q_a, q_c, e_s_result, rag_output = query_processor.muti_linkup(
            query_list,
            top_k=request.top_k,
            modes=request.modes
        )
        
        # 格式化返回结果
        return SearchResult(
            qc_results=[
                {
                    "title": item["title"],
                    "content": item["content"],
                    "score": item["score"]
                }
                for item in q_c[0]  # 假设 q_c 是列表的列表
            ],
            es_results=[
                {
                    "title": item["title"],
                    "content": item["content"],
                    "score": item["score"]
                }
                for item in e_s_result[0]  # 假设 e_s_result 是列表的列表
            ],
            gpt_answer=rag_output[0]  # 假设 rag_output 是答案列表
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
